# 文件: upload_service.py
import os

def save_chunk(key, file_id, chunk_index, data, temp_folder):
    """
    保存分片数据到临时目录。

    :param key: 分片对应的属性名
    :param file_id: 文件标识符
    :param chunk_index: 当前分片的索引
    :param data: 分片数据
    :param temp_folder: 临时存储分片的根目录
    """
    temp_key_folder = os.path.join(temp_folder, file_id, key)
    os.makedirs(temp_key_folder, exist_ok=True)

    temp_file_path = os.path.join(temp_key_folder, f"chunk_{chunk_index}")
    if isinstance(data, str):  # 如果是字符串数据（非 Blob 类型）
        with open(temp_file_path, "w") as f:
            f.write(data)  # 将字符串写入文件
    else:  # 如果是 Blob 数据
        data.save(temp_file_path)  # 保存为二进制文件

    return temp_key_folder

def merge_chunks(key, file_id, total_chunks, temp_folder, upload_folder):
    """
    合并所有分片数据为完整文件。

    :param key: 属性名
    :param file_id: 文件标识符
    :param total_chunks: 总分片数量
    :param temp_folder: 临时存储分片的根目录
    :param upload_folder: 最终文件存储目录
    """
    temp_key_folder = os.path.join(temp_folder, file_id, key)
    final_file_path = os.path.join(upload_folder, f"{file_id}_{key}")

    with open(final_file_path, "wb") as final_file:
        for i in range(total_chunks):
            chunk_path = os.path.join(temp_key_folder, f"chunk_{i}")
            with open(chunk_path, "rb") as chunk_file:
                final_file.write(chunk_file.read())  # 将每个分片内容写入最终文件
            os.remove(chunk_path)  # 删除临时分片

    os.rmdir(temp_key_folder)  # 删除临时目录
    return final_file_path

def is_all_chunks_uploaded(temp_key_folder, total_chunks):
    """
    检查是否接收到了所有分片。

    :param temp_key_folder: 临时目录路径
    :param total_chunks: 总分片数量
    :return: 是否所有分片已上传
    """
    uploaded_chunks = len([f for f in os.listdir(temp_key_folder) if f.startswith("chunk_")])
    return uploaded_chunks == total_chunks


# 文件: app.py
from flask import Flask, request, jsonify
from upload_service import save_chunk, merge_chunks, is_all_chunks_uploaded
import os

app = Flask(__name__)

# 文件上传主目录
UPLOAD_FOLDER = "uploads"
# 临时存储分片的目录
TEMP_FOLDER = os.path.join(UPLOAD_FOLDER, "temp")
os.makedirs(TEMP_FOLDER, exist_ok=True)

@app.route("/upload", methods=["POST"])
def upload_chunk():
    """
    接收客户端上传的分片数据并存储，
    在所有分片上传完成后合并成完整的文件。
    """
    # 从请求中提取必要的元数据
    key = request.form.get("key")  # 当前分片对应的属性名
    chunk_index = int(request.form.get("chunk_index", -1))  # 当前分片的索引
    total_chunks = int(request.form.get("total_chunks", -1))  # 总分片数量
    data = request.files.get("data") or request.form.get("data")  # 分片数据

    if not key or chunk_index == -1 or total_chunks == -1 or not data:
        return jsonify({"error": "Invalid request"}), 400

    # 文件标识符（例如通过用户会话或文件 ID 区分不同文件）
    file_id = request.form.get("file_id", "default_file")

    # 保存当前分片
    temp_key_folder = save_chunk(key, file_id, chunk_index, data, TEMP_FOLDER)

    # 检查是否所有分片已上传
    if is_all_chunks_uploaded(temp_key_folder, total_chunks):
        # 合并所有分片
        merge_chunks(key, file_id, total_chunks, TEMP_FOLDER, UPLOAD_FOLDER)
        return jsonify({"message": f"Key '{key}' uploaded successfully"}), 200

    return jsonify({"message": f"Chunk {chunk_index + 1}/{total_chunks} for key '{key}' uploaded"}), 200

if __name__ == "__main__":
    # 启动 Flask 应用
    app.run(debug=True)
